import os
import sys

def generate_pub_mod(src_dir):

    content = ''
    if os.path.isdir(src_dir):
        content = 'pub mod {}'.format(os.path.basename(src_dir)) + '{\n'
        for sub_dir in os.listdir(src_dir):
            content += generate_pub_mod(os.path.join(src_dir, sub_dir))
        content += '}\n'
    else:
        filename = str(os.path.basename(src_dir)).split('.')[0]
        content = f'pub mod {filename};\n'

    return content

def main(rust_root_dir):
    # 构建lib.rs文件的路径
    lib_rs_path = os.path.join(rust_root_dir, 'lib.rs')
    os.remove(lib_rs_path)
    pub_mod_content = generate_pub_mod(os.path.join(rust_root_dir, 'src'))
    file_content = ("#![allow(dead_code)]\n"
                    "#![allow(mutable_transmutes)]\n"
                    "#![allow(non_camel_case_types)]\n"
                    "#![allow(non_snake_case)]\n"
                    "#![allow(non_upper_case_globals)]\n"
                    "#![allow(unused_assignments)]\n"
                    "#![allow(unused_mut)]\n"
                    "#![feature(label_break_value)]\n"
                    "#![feature(extern_types)]\n"
                    "#![feature(adt_const_params)]\n"
                    "extern crate libc;\n")
    file_content += pub_mod_content
    with open(lib_rs_path, 'w') as f:
        f.write(file_content)


if __name__ == '__main__':
    main(sys.argv[1])
